package alias_sample

import (
	"container/list"
	"math/rand"
)
type AliasSample struct {
    accept []float64 // 属于自己本身的概率
    alias  []int     // 上半部分从哪里来的
}

// 构建采样器，核心是构建数组accept和alias
func NewAliasSample(prob []float64) *AliasSample {
    if len(prob) == 0 {
        return nil
    }
    // 先将概率归1化
    sum := 0.0
    for _, v := range prob {
        sum += v
    }
    for i := range prob {
        prob[i] /= sum
    }
    // 每个事件乘n
    n := float64(len(prob))
    arr := make([]float64, len(prob))
    for i := range arr {
        arr[i] = n * prob[i]
    }
    // 初始化accept 和 alias
    accept := make([]float64, len(prob))
    alias := make([]int, len(prob))
    for i := 0; i < len(prob); i++ {
        accept[i] = 1.0
        alias[i] = -1
    }
    // 概率小于1的放入到small栈置，概率大于等于1的放入到large
    small := list.New()
    large := list.New()
    for i, v := range arr {
        if v < 1.0 {
            small.PushBack(i)
        } else {
            large.PushBack(i)
        }
    }
    // 将面积大于一的多余的面积补充到小面积中，确保每个面积都为1
    for small.Len() > 0 && large.Len() > 0 {
        // 分别从栈中取出两个元素
        smallIndex := small.Remove(small.Front()).(int)
        largeIndex := large.Remove(large.Front()).(int)
        // 属于自己的概率
        accept[smallIndex] = arr[smallIndex]
        // 上半部分从哪里来
        alias[smallIndex] = largeIndex
        // 剩余的面积
        arr[largeIndex] -= (1 - arr[smallIndex])
        // 如果剩余的面积大于1，则放入large栈，否则放入small栈
        if arr[largeIndex] >= 1.0 {
            large.PushBack(largeIndex)
        } else {
            small.PushBack(largeIndex)
        }
    }
    return &AliasSample{accept, alias}
}

// 采样，支持并发调用
func (s *AliasSample) Sample() int {
    i := rand.Intn(len(s.accept)) // 随机选择一个元素
	f := rand.Float64() // 随机生成一个概率
	if f < s.accept[i] {
		return i // 属于自己
	}
	return s.alias[i] // 上半部分从哪里来
}